import numpy as np

class QuinticSample:# 5
    def __init__(self, origion, end, var_array):
        self.origion = origion
        self.end = end
        self.var_array = var_array

    def run(self):
        # p_range = self.end["var"]
        # v_range = self.end["v"]
        # a_range = self.end["a"]
        #
        # p_array = np.arange(p_range[0], p_range[1], self.nums)
        # v_array = np.arange(v_range[0], v_range[1], self.nums)
        # a_array = np.arange(a_range[0], a_range[1], self.nums)
        # finalstate_array = []
        # for i in range(self.nums):
        #     finalstate_array.append([p_array[i], v_array[i], a_array[i]])
        coeff_list = []
        for t_index in range(len(self.var_array)):
            t=self.var_array[t_index]
            for state_index in range(len(self.end["var"])):
                final_state=[self.end["var"][state_index],self.end["var_dot"][state_index],self.end["var_2dot"][state_index]]
                O = self.origion + final_state
                M = np.array([
                    [1, 0, 0, 0, 0, 0],
                    [0, 1, 0, 0, 0, 0],
                    [0, 0, 2, 0, 0, 0],
                    [1, t, t ** 2, t ** 3, t ** 4, t ** 5],
                    [0, 1, 2 * t, 3 * t ** 2, 4 * t ** 3, 5 * t ** 4],
                    [0, 0, 2, 6 * t, 12 * t ** 2, 20 * t ** 3]
                ])
                O = np.array([O]).T
                t_co = np.dot(np.linalg.inv(M), O)
                t_co = np.flipud(t_co).flatten()
                t_fun = np.poly1d(t_co.astype(float))
                coeff_list.append({"var": t, "f": t_fun})
        return coeff_list



class QuarticSample:# 4
    def __init__(self, origion, end, steptimes):
        self.origion = origion
        self.end = end
        self.steptimes = steptimes

    def run(self):
        coeff_list = []
        for t_index in range(len(self.steptimes)):
            t=self.steptimes[t_index]
            for state_index in range(len(self.end["v"])):
                O = [0] + self.origion + [self.end["v"][state_index],self.end["a"][state_index]]
                M = [
                    [1, 0, 0, 0, 0],  # 初始位移
                    [0, 1, 0, 0, 0],  # 初始速度
                    [0, 0, 2, 0, 0],  # 初始加速度
                    [0, 1, 2 * t, 3 * t ** 2, 4 * t ** 3],
                    [0, 0, 2, 6 * t, 12 * t ** 2]
                ]
                O = np.array([O]).T
                t_co = np.dot(np.linalg.inv(M), O)
                t_co = np.flipud(t_co).flatten()
                t_fun = np.poly1d(t_co.astype(float))
                coeff_list.append({"var": t, "f": t_fun})
        return coeff_list